#include "../include/ForceControl.hpp"

#include <cmath>
#include <iostream>
#include <string>


typedef std::chrono::high_resolution_clock Clock;

using namespace dynamica;
using std::string;
using std::cout;
using std::endl;

ForceControlController::ForceControlController()
{
    _pose_user_input      = new double[7]();
    _wrench_Tr_set        = Vector6::Zero();
    _Tr                   = Matrix6::Identity();
    _Tr_inv               = Matrix6::Identity();
    _m_force_selection    = Matrix6::Zero();
    _m_velocity_selection = Matrix6::Identity();

    _pose_sent_to_robot = new double[7]();
    _v_W                = Vector6::Zero();
    _v_T                = Vector6::Zero();
    _wrench_T_Err      = Vector6::Zero();
    _wrench_T_Err_I    = Vector6::Zero();
    _SE3_WT_old         = Matrix4::Identity();
    _SE3_WToffset       = Matrix4::Identity();
}

ForceControlController::~ForceControlController()
{
    delete [] _pose_user_input;
    delete [] _pose_sent_to_robot;

    if (_print_flag)
        _file.close();
}

bool ForceControlController::init_controller(){
    cout << "[ForceControlController] Begin initialization.\n";
        /* use current state to initialize commands */
    double wrench[6];
    _hw->getState(_pose_user_input, wrench);
    stream_array_in(cout, _pose_user_input, 7);
    cout << endl;
    copyArray(_pose_user_input, _pose_sent_to_robot, 7);
    _SE3_WT_old = posemm2SE3(_pose_user_input);

    /**
     * Force control parameters
     */
    _FC_I_limit_T_6D = Vector6(FC_I_Limit_T_6D_elements.data());
    /**
     * Spring-mass-damper coefficients
     */
    _dt = 1.0/_fHz;
    _ToolStiffnessMatrix = Vector6(Stiffness_matrix_diag_elements.data()).asDiagonal();
    _ToolDamping_coef    = Vector6(Damping_matrix_diag_elements.data()).asDiagonal();
    _ToolInertiaMatrix   = Vector6(Inertia_matrix_diag_elements.data()).asDiagonal();

    if (_print_flag) {
        _file.open(fullpath);
        if (_file.is_open())
            printf("[ForceControlController] file opened successfully.\n");
        else
            printf("[ForceControlController] Failed to open file.\n");
    }
    cout << "[ForceControlController] initialization is done." << endl;
    return true;
}

bool ForceControlController::init(ForceControlHardware* hw, std::chrono::high_resolution_clock::time_point time0)
{
    _hw    = hw;
    _time0 = time0;
    return true;
}

void ForceControlController::setPose(const double *pose)
{
    copyArray(pose, _pose_user_input, 7);
}

void ForceControlController::setForce(const double *force)
{
    _wrench_Tr_set(0) = force[0];
    _wrench_Tr_set(1) = force[1];
    _wrench_Tr_set(2) = force[2];
    _wrench_Tr_set(3) = force[3];
    _wrench_Tr_set(4) = force[4];
    _wrench_Tr_set(5) = force[5];
}

void ForceControlController::getPose(double *pose) {
    _hw->getPose(pose);
}

void ForceControlController::getToolVelocity(Eigen::Matrix<double, 6, 1> *v_T) {
    *v_T = _v_T;
}

bool ForceControlController::getToolWrench(Eigen::Matrix<double, 6, 1> *wrench) {
    double wrench_array[6];
    bool safety = _hw->getWrench(wrench_array);
    for(int i = 0; i < 6; i++)
        (*wrench)[i] = wrench_array[i];
    return safety;
}



/*
 *
    force control law
        Frames/spaces:
            W: world frame
            T: current tool frame
            So: set tool frame with offset
            Tf: transformed generalized space
        Quantities:
            SE3: 4x4 homogeneous coordinates
            se3: 6x1 twist coordinate of SE3
            spt: 6x1 special twist: 3x1 position, 3x1 exponential coordinate for rotation
            td: 6x1 time derivative of twist.
            v: 6x1 velocity, either spatial or body
            wrench: 6x1 wrench. Makes work with body velocity
 *
 */
int ForceControlController::update() {
    // _hw->setPose(_pose_user_input);
    // return 0;
    double pose_fb[7];
    double wrench_fb[6];
    // 0: no error. 1: still waiting for new data. 2: dead stream.
    //      *             3: force is too big
    int force_feedback_code;
    while (true) {
        force_feedback_code = _hw->getState(pose_fb, wrench_fb);
        if (force_feedback_code == 0) {
            // no error, use this data
            break;
        } else {
            // dangerous, stop execution
            return force_feedback_code;
        }
    }
    // printf("F fb:  %.3f\t%.3f\t%.3f\t%.3f\t%.3f\t%.3f\n", 
    //     wrench_fb[0],
    //     wrench_fb[1],
    //     wrench_fb[2],
    //     wrench_fb[3],
    //     wrench_fb[4],
    //     wrench_fb[5]);

    // ----------------------------------------
    //  Compute Forces in Generalized space
    // ----------------------------------------
    /*  Spring forces  (using stiffness)   */
    // Pose Error
    Matrix4 SE3_WTSet;
    SE3_WTSet = posem2SE3(_pose_user_input);

    Matrix4 SE3_WSo;
    SE3_WSo = _SE3_WToffset * SE3_WTSet;

    Matrix4 SE3_WT_fb, SE3_TSo;
    //SE3
    SE3_WT_fb = posem2SE3(pose_fb);
    
    SE3_TSo   = SE3Inv(SE3_WT_fb)*SE3_WSo; // aka SE3_S_err

    Vector6 spt_TSo;
    spt_TSo = SE32spt(SE3_TSo);

    // Jac0 * body velocity = spt time derivative
    Matrix6 Jac0, Jac0_inv;
    Jac0_inv = JacobianSpt2BodyV(SE3_WT_fb.block<3,3>(0,0));
    Jac0 = Jac0_inv.inverse();

    // elastic wrench
    Vector6 wrench_T_spring;
    wrench_T_spring = Jac0*_ToolStiffnessMatrix*spt_TSo;

    /*  Force feedback  */
    // note: the minus sign here make the physical meaning of the force correct:
    //      the force being acted on the environment from the robot
    Vector6 wrench_T_fb; // force feedback measured in tool frame
    for (int i = 0; i < 6; ++i) wrench_T_fb(i) = - wrench_fb[i];

    /* velocity */
    Matrix6 Adj_WT, Adj_TW;
    Adj_WT = SE32Adj(SE3_WT_fb);
    Adj_TW = SE32Adj(SE3Inv(SE3_WT_fb));
    _v_T  =  Adj_TW * _v_W;
    Vector6 v_Tr;
    v_Tr =  _Tr * _v_T;
    
    /* transformation from Tool wrench to
            transformed space  */
    Vector6 wrench_Tr_spring, wrench_Tr_fb;
    wrench_Tr_spring = _Tr*wrench_T_spring;

    /*  Force error, PID force control */
    Vector6 wrench_T_Set, wrench_T_Err;
    wrench_T_Set = _Tr_inv*_wrench_Tr_set;

    // std::cout << "wrench_T_Set \n" << wrench_T_Set << "\n"; 
    
    wrench_T_Err = wrench_T_Set - wrench_T_fb;
    _wrench_T_Err_I += wrench_T_Err;
    truncate6d(&_wrench_T_Err_I, -_FC_I_limit_T_6D, _FC_I_limit_T_6D);

    Vector6 wrench_T_PID;
    wrench_T_PID.head(3) = _kForceControlPGainTran*wrench_T_Err.head(3)
            + _kForceControlIGainTran*_wrench_T_Err_I.head(3)
            + _kForceControlDGainTran*(wrench_T_Err.head(3) - _wrench_T_Err.head(3));
    wrench_T_PID.tail(3) = _kForceControlPGainRot*wrench_T_Err.tail(3)
            + _kForceControlIGainRot*_wrench_T_Err_I.tail(3)
            + _kForceControlDGainRot*(wrench_T_Err.tail(3) - _wrench_T_Err.tail(3));
    Vector6 wrench_Tr_PID;
    wrench_Tr_PID = _Tr * wrench_T_PID;
    _wrench_T_Err = wrench_T_Err;
    Vector6 wrench_Tr_Err;
    wrench_Tr_Err = _Tr*wrench_T_Err;

    Vector6 wrench_Tr_damping;
    wrench_Tr_damping = - _Tr*_ToolDamping_coef*_v_T;

    Vector6 wrench_Tr_All;
    wrench_Tr_All = _m_force_selection *
            (wrench_Tr_spring + wrench_Tr_Err + wrench_Tr_PID + wrench_Tr_damping);

    // ----------------------------------------
    //  force to velocity
    // ----------------------------------------

    /* Newton's Law */
    //  Axes are no longer independent when we take
    //      rotation in to consideration.
    //  Newton's Law in body (Tool) frame:
    //      W=M*vd
    //          W: body wrench
    //          M: Inertia matrix in body frame
    //          vd: body velocity time derivative
    //  Newton's law in transformed space
    //      TW=TMTinv Tvd
    //      W_Tr = TMTinv vd_Tr
    Matrix6 Tinv;
    Tinv = _Tr.inverse();
    Vector6 vd_Tr;
    vd_Tr = (_Tr*_ToolInertiaMatrix*Tinv).fullPivLu().solve(wrench_Tr_All);

    // integration
    // right now the velocity vector _vd_Tr only contains force command
    v_Tr += _dt * vd_Tr;
    v_Tr = _m_force_selection * v_Tr;

    /* Velocity command */
    Vector6 v_T_command;
    v_T_command = Jac0_inv * spt_TSo / _dt;
    v_Tr += _m_velocity_selection*_Tr*v_T_command;
    _v_W = Adj_WT * Tinv * v_Tr;

    // ----------------------------------------
    //  velocity to pose
    // ----------------------------------------
    Matrix4 SE3_WT_command;
    SE3_WT_command = SE3_WT_fb + wedge6(_v_W)*SE3_WT_fb*_dt;
    SE32Posem(SE3_WT_command, _pose_sent_to_robot);

    Clock::time_point timenow_clock = Clock::now();
    double timenow = double(std::chrono::duration_cast<std::chrono::nanoseconds>(
            timenow_clock - _time0).count())/1e6; // milli second

    // if (std::isnan(_pose_sent_to_robot[0])) {
        // cout << "==================== Temp variables: =====================\n";
        // cout << "pose_fb: ";
        // stream_array_in(cout, pose_fb, 7);
        // cout << "\n_pose_sent_to_robot: ";
        // stream_array_in(cout, _pose_sent_to_robot, 7);
        // cout << "\nwrench_Tr_spring: \n" << wrench_Tr_spring.format(MatlabFmt) << endl;
        // // cout << "wrench_T_fb: \n" << wrench_T_fb.format(MatlabFmt) << endl;
        // // cout << "wrench_Tr_fb: \n" << wrench_Tr_fb.format(MatlabFmt) << endl;
        // // cout << "_wrench_Tr_set: \n" << _wrench_Tr_set.format(MatlabFmt) << endl;
        // cout << "wrench_T_PID: \n" << wrench_T_PID.format(MatlabFmt) << endl;
        // cout << "wrench_Tr_damping: \n" << wrench_Tr_damping.format(MatlabFmt) << endl;
        // cout << "wrench_Tr_Err: \n" << wrench_Tr_Err.format(MatlabFmt) << endl;
        // cout << "wrench_Tr_All: \n" << wrench_Tr_All.format(MatlabFmt) << endl;
        // cout << "vd_Tr: \n" << vd_Tr.format(MatlabFmt) << endl;
        // cout << "v_Tr: \n" << v_Tr.format(MatlabFmt) << endl;
        // cout << "v_T_command: \n" << v_T_command.format(MatlabFmt) << endl;
        // cout << "_v_W: \n" << _v_W.format(MatlabFmt) << endl;
        // cout << "SE3_WT_fb: \n" << SE3_WT_fb.format(MatlabFmt) << endl;
        // cout << "SE3_WT_command: \n" << SE3_WT_command.format(MatlabFmt) << endl;
    //     displayStates();
    //     cout << "Press ENTER to continue..." << endl;
    //     getchar();
    // }

    //lock y 
    _pose_sent_to_robot[1]=0;
    _hw->setPose(_pose_sent_to_robot);

    // cout << "[ForceControlController] Update at "  << timenow << endl;
    if(_print_flag)
    {
        _file << timenow << " ";
        stream_array_in(_file, _pose_user_input, 7);
        stream_array_in(_file, pose_fb, 7);
        stream_array_in6d(_file, wrench_T_Set);
        stream_array_in(_file, wrench_T_fb, 6);
        stream_array_in6d(_file, wrench_Tr_All);
        stream_array_in(_file, _pose_sent_to_robot, 7);
        _file << endl;
    }
    return force_feedback_code;
}

// After axis update, the goal pose with offset should not have error in
// velocity controlled axes of the new axes. To satisfy this requirement,
// we need to change _SE3_WToffset accordingly
void ForceControlController::updateAxis(const Matrix6 &Tr, int n_af)
{
    Matrix4 SE3_WTSet;
    SE3_WTSet = posem2SE3(_pose_user_input);//set user input

    Matrix4 SE3_WSo;
    SE3_WSo = _SE3_WToffset * SE3_WTSet;

    Matrix4 SE3_WT_fb, SE3_TSo;
    SE3_WT_fb = posem2SE3(_pose_sent_to_robot);
    SE3_TSo   = SE3Inv(SE3_WT_fb)*SE3_WSo; // aka SE3_S_err

    // std::cout << "SE3_WSo\n" << SE3_WSo << "\n";
    // std::cout << "SE3Inv(SE3_WT_fb)\n" << SE3Inv(SE3_WT_fb) << "\n";

    Vector6 spt_TSo;
    spt_TSo = SE32spt(SE3_TSo);

    Vector6 v_force_selection, v_velocity_selection;
    v_force_selection << 0, 0, 0, 0, 0, 0;
    v_velocity_selection << 1, 1, 1, 1, 1, 1;
    for (int i = 0; i < n_af; ++i) {
        v_force_selection(i) = 1;
        v_velocity_selection(i) = 0;
    }

    _m_force_selection    = v_force_selection.asDiagonal();
    _m_velocity_selection = v_velocity_selection.asDiagonal();

    Matrix6 Jac0, Jac0_inv;
    Jac0_inv = JacobianSpt2BodyV(SE3_WT_fb.block<3, 3>(0, 0));
    Jac0 = Jac0_inv.inverse();

    // std::cout << "jac0\n" <<Jac0 << "\njac0_inv\n" << Jac0_inv << "\n";  

    MatrixXd m_anni = _m_velocity_selection*Tr*Jac0;

    // std::cout << "m_anni\n" << m_anni << "\n";
    Vector6 spt_TSo_new = (Matrix6::Identity() - pseudoInverse(m_anni, 1e-6)*m_anni)*spt_TSo;

    // std::cout << "spt_TSo_new\n" << spt_TSo_new << "\n";
    // update offset
    _SE3_WToffset = SE3_WT_fb*spt2SE3(spt_TSo_new)*SE3Inv(SE3_WTSet);

    // std::cout << "_SE3_WToffset\n" << _SE3_WToffset << "\n";

    // project these into force space
    _wrench_T_Err_I = _Tr_inv*_m_force_selection*Tr*_wrench_T_Err_I;
    _wrench_T_Err   = _Tr_inv*_m_force_selection*Tr*_wrench_T_Err;

    _SE3_WT_old = SE3_WT_fb;
    _Tr         = Tr;
    _Tr_inv     = _Tr.inverse();

    if (std::isnan(_SE3_WToffset(0,0))) {
        // cout << "SE3_WT_fb:\n" << SE3_WT_fb.format(MatlabFmt) << endl;
        // cout << "SE3_TSo:\n" << SE3_TSo.format(MatlabFmt) << endl;
        // cout << "spt_TSo:\n" << spt_TSo.format(MatlabFmt) << endl;
        // cout << "Jac0_inv:\n" << Jac0_inv.format(MatlabFmt) << endl;
        // cout << "Jac0:\n" << Jac0.format(MatlabFmt) << endl;
        // cout << "m_anni:\n" << m_anni.format(MatlabFmt) << endl;
        // cout << "spt_TSo_new:\n" << spt_TSo_new.format(MatlabFmt) << endl;
        // cout << "_SE3_WToffset:\n" << _SE3_WToffset.format(MatlabFmt) << endl;
        cout << "\nNow paused at updateAxis()";
        getchar();
    }
}

bool ForceControlController::ExecuteHFVC(const int n_af, const int n_av,
  const Matrix6 &R_a, const double *pose_set, const double *force_set,
  HYBRID_SERVO_MODE mode, const int main_loop_rate, const double duration) {
    assert(n_af + n_av == 6);
    if (mode == HS_STOP_AND_GO) {
      reset();
    }
    updateAxis(R_a, n_af);
    setForce(force_set);

    // get current pose for motion planning
    double pose_fb[7];
    if (mode == HS_STOP_AND_GO)
      _hw->getPose(pose_fb);
    else
      copyArray(_pose_user_input, pose_fb, 7);

    /* Motion Planning */
    // get num of step in integer
    int num_of_steps = round(double(main_loop_rate) * duration); 
    MatrixXd pose_traj;
    MotionPlanningLinear(pose_fb, pose_set, num_of_steps, &pose_traj);
    // MotionPlanningTrapezodial(pose_fb, pose_set, _kAccMaxTrans, _kVelMaxTrans,
    //         _kAccMaxRot, _kVelMaxRot, (double)main_loop_rate, &pose_traj);

    /* Execute the motion plan */
    // ros::Rate pub_rate(main_loop_rate);
    int b_unsafe = false;
    for (int i = 0; i < num_of_steps; ++i) {
      // cout << "[Hybrid] update step " << i << " of " << num_of_steps;
      // cout << ", pose sent: " << pose_traj(0, i) << ", " << pose_traj(1, i);
      // cout << ", " << pose_traj(2, i) << endl;
      setPose(pose_traj.col(i).data());
      // !! after setPose, must call update before updateAxis
      // so as to set correct value for pose_command
      b_unsafe = update();
      if(b_unsafe != 0) {
        // ROS_ERROR_STREAM("[force_control] Unsafe force feedback! error code: " << b_unsafe);
        break;
    }
        Sleep(_control_delta_time);
    //   pub_rate.sleep();
    }
    return b_unsafe;
}

bool ForceControlController::ExecuteLinearPos(const double *pose_set, const int main_loop_rate, const double duration){
    double pose_fb[7];
    _hw->getPose(pose_fb);
    int num_of_steps = round(double(main_loop_rate) * duration); 
    MatrixXd pose_traj;
    MotionPlanningLinear(pose_fb, pose_set, num_of_steps, &pose_traj);
    for (int i = 0; i < num_of_steps; ++i) {
        setPose(pose_traj.col(i).data());
        _hw->setPose(_pose_user_input);

        for(auto i=0; i<3; i++){
            printf("_pose_user_input \t %.4f\t%.4f\t%.4f\n", _pose_user_input[0], _pose_user_input[1], _pose_user_input[2]);
        }
        Sleep(_control_delta_time);
    }
    return true;
}

// reset everytime you start from complete stop.
// clear up internal states
void ForceControlController::reset()
{
    _hw->getPose(_pose_sent_to_robot);
    _hw->getPose(_pose_user_input);
    _SE3_WT_old = posem2SE3(_pose_sent_to_robot);
    _SE3_WToffset = Matrix4::Identity();
    _v_W = Vector6::Zero();
    _v_T = Vector6::Zero();
    _wrench_T_Err = Vector6::Zero();
    _wrench_T_Err_I = Vector6::Zero();
}

void ForceControlController::displayStates() {
    using namespace std;
    // cout << "================= Parameters ================== " << endl;
    // cout << "_ToolStiffnessMatrix: \n" << _ToolStiffnessMatrix.format(MatlabFmt)
    //         << endl;
    // cout << "_ToolDamping_coef: \n" << _ToolDamping_coef.format(MatlabFmt)
    //         << endl;
    // cout << "_ToolInertiaMatrix: \n" << _ToolInertiaMatrix.format(MatlabFmt)
    //         << endl;
    // cout << "================= Commands ================== " << endl;
    // cout << "_pose_user_input: ";
    // stream_array_in(cout, _pose_user_input, 7);
    // cout << "\n_Wrench_Tr_set: \n" << _wrench_Tr_set.format(MatlabFmt) << endl;
    // cout << "_Tr: \n" << _Tr.format(MatlabFmt) << endl;
    // cout << "_Tr_inv: \n" << _Tr_inv.format(MatlabFmt) << endl;
    // cout << "_m_force_selection: \n" << _m_force_selection.format(MatlabFmt) << endl;
    // cout << "_m_velocity_selection: \n" << _m_velocity_selection.format(MatlabFmt) << endl;
    // cout << "================= Internal states ================== " << endl;
    // cout << "_pose_sent_to_robot: ";
    // stream_array_in(cout, _pose_sent_to_robot, 7);
    // cout << "\n_SE3_WT_old: \n" << _SE3_WT_old.format(MatlabFmt) << endl;
    // cout << "_SE3_WToffset: \n" << _SE3_WToffset.format(MatlabFmt) << endl;
    // cout << "_v_W: \n" << _v_W.format(MatlabFmt) << endl;
    // cout << "_wrench_T_Err: \n" << _wrench_T_Err.format(MatlabFmt) << endl;
    // cout << "_wrench_T_Err_I: \n" << _wrench_T_Err_I.format(MatlabFmt) << endl;
}